package com.hys.ai.controller;

import cn.hutool.core.util.IdUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * @author hys
 * @version Id: TongyiSimpleController.java, v 0.1 2024年07月31日 09:22 hys Exp $
 */
@Slf4j
@RestController
public class DemoController {

    @Autowired
    private       OllamaChatModel chatModel;
    // 模拟数据库存储会话和消息
    private final ChatMemory      chatMemory = new InMemoryChatMemory();

    @RequestMapping(value = "/demo/call")
    public String call(@RequestParam(defaultValue = "你是") String prompt) {
        log.info("=========<" + prompt);
        String content = chatModel.call(prompt);
        log.info("=========>" + content);
        return content;
    }

    @RequestMapping(value = "/demo/stream", produces = "text/event-stream; charset=utf-8")
    public Flux stream(@RequestBody RequestObject request) {
        //String requestId = IdUtil.simpleUUID();
        //System.out.println(request.getPrompt());
        Flux<ServerSentEvent<SSEData>> message = chatModel.stream(request.getPrompt()).map(resp -> {
            SSEData sseData = new SSEData(resp);
            return ServerSentEvent.<SSEData>builder()
                    .data(sseData)
                    .build();
        }).concatWith(Mono.just(createDoneEvent(null)));
        return message;
    }

    private ServerSentEvent<SSEData> createDoneEvent(String requestId) {
        SSEData doneData = new SSEData("[DONE]");
        return ServerSentEvent.<SSEData>builder()
                .data(doneData)
                .build();
    }

}
